import argparse
import io
import os
import shutil
import sys
from datetime import datetime


import rembg
from PIL import Image
from rembg import new_session

from crop import crop_image
from pad import resize_and_pad_image


# pyinstaller -F main.py -n matting;pyinstaller .\matting.spec

def remove_background(input_path, output_path, session=None, alpha_matte=True, background_color=(255, 255, 255, 0)):
    with open(input_path, "rb") as input_file:
        input_data = input_file.read()
        # Use advanced options
        output_data = rembg.remove(input_data, session=session, alpha_matte=alpha_matte, bgcolor=background_color)
        image_stream = io.BytesIO(output_data)
        img = Image.open(image_stream)
        # output_file.write(output_data)
        img = crop_image(img)
        img = img.convert('RGB')  # PNG 通常是 RGBA，需要转换为 RGB
        img = resize_and_pad_image(img)
        img.save(output_path, 'JPEG')



def check_image_format(file_path):
    try:
        with open(file_path, 'rb') as f:
            header = f.read(8)  # Read the first 8 bytes of the file

            if header.startswith(b'\xff\xd8'):
                return 'JPEG'
            elif header.startswith(b'\x89PNG\r\n\x1a\n'):
                return 'PNG'
            else:
                return 'Unknown format'
    except IOError:
        return 'File not found or unreadable'

def get_filelist(current_directory=None):
    # 获取当前目录路径
    if current_directory is None:
        current_directory = os.getcwd()

    # 列出当前目录下的所有文件和子目录
    filelist = os.listdir(current_directory)

    # 过滤出文件
    files = [f for f in filelist if os.path.isfile(os.path.join(current_directory, f))]

    return files


def get_current_date():
    # 获取当前日期
    today = datetime.now()
    # 格式化日期为 '年-月-日'
    formatted_date = today.strftime('%Y-%m-%d')
    return formatted_date

def print_progress_bar(iteration, total, length=50, message=''):
    percent = ("{0:.1f}").format(100 * (iteration / float(total)))
    filled_length = int(length * iteration // total)
    bar = '█' * filled_length + '-' * (length - filled_length)
    sys.stdout.write(f'\r[{bar}] {percent}% Complete - {message}')
    sys.stdout.flush()

def matting_parser():
    parser = argparse.ArgumentParser(description="批处理抠图软件")
    parser.add_argument('-i', '--input', type=str, help='[可选]输入文件路径，默认当前路径')
    parser.add_argument('-o', '--output', type=str, help='[可选]输出文件路径，默认以当前日期创建在当前目录')
    parser.add_argument('-m', '--model', type=str, help='[可选]自定义模型')
    return parser

working_directory = "./"
image_extension = ['jpeg', 'jpg', 'png']
model_path = './u2net.onnx'
if __name__ == '__main__':
    try:
        # 解析参数
        parser = matting_parser()
        args = parser.parse_args()

        # create output directory
        output_dir = get_current_date()
        os.makedirs(output_dir, exist_ok=True)
        output_done_dir = output_dir + '-orig'
        os.makedirs(output_done_dir, exist_ok=True)


        # 打印解析后的参数
        if args.input:
            working_directory = args.input
        if args.output:
            output_dir = args.output
        if args.model:
            model_path = args.model

        print(f'输入目录: {working_directory}')
        print(f'输出目录: {output_dir}')

        image_list = get_filelist(working_directory)
        total = len(image_list)
        if len(image_list) != 0:
            iter = 0
            session = new_session("u2net_custom", model_path=model_path)
            print("模型加载成功!")
            for file in image_list:
                image = os.path.basename(file)
                filename, ext = os.path.splitext(image)
                if ext.lstrip('.').lower() in image_extension:
                    output_file = os.path.join(output_dir, filename + ".jpg")
                    remove_background(image, output_file, session=session)
                    shutil.move(image, output_done_dir)
                    # print(f"{output_file} is {check_image_format(output_file)}")
                iter += 1
                print_progress_bar(iter, total, message=image)
            print("\n处理完成!")
        else:
            print("没有需要处理的图片")

    except Exception as e:
        print(f"An error occurred: {e}")
